import numpy as np



# ------------------------------BilinearInterpolator--------------------------------
def csharp_binary_search(arr, key):
    """
    :param arr: 顺序数组
    :param key: 要查找的key
    :return: https://learn.microsoft.com/zh-cn/dotnet/api/system.array.binarysearch?view=net-8.0
    """
    left, right = 0, len(arr)-1
    while left <= right:
        mid = (left + right) >> 1
        if arr[mid] == key:
            return mid
        if arr[mid] < key:
            left = mid + 1
        else:
            right = mid - 1
    return ~left


class xh_BI:
    def __init__(self, x:np.ndarray, y:np.ndarray, z:np.ndarray):
        """
        :param x: 一维
        :param y: 一维
        :param z: 二维
        """
        self.x_arr = x
        self.y_arr = y
        self.z_arr = z

    def interpolate(self, x:np.ndarray, y:np.ndarray):
        x_i = csharp_binary_search(self.x_arr, x)
        if x_i < 0:
            x_i = ~x_i - 1
        x_i = np.clip(x_i, 0, self.x_arr.size - 2)
        y_j = csharp_binary_search(self.y_arr, y)
        if y_j < 0:
            y_j = ~y_j - 1
        y_j = np.clip(y_j, 0, self.y_arr.size - 2)

        x0, x1 = self.x_arr[x_i], self.x_arr[x_i + 1]
        y0, y1 = self.y_arr[y_j], self.y_arr[y_j + 1]
        # ----------------------------------------
        # 数组以左上角为原点，y轴变横轴，x轴变纵轴
        # ----------------------------------------
        # z00, z01 = self.z_arr[x_i,        y_j], self.z_arr[x_i,        y_j + 1]
        # z10, z11 = self.z_arr[x_i + 1, y_j], self.z_arr[x_i + 1, y_j + 1]
        # -----------------------
        # 修正二：
        # -----------------------
        z00, z01 = self.z_arr[y_j,        x_i], self.z_arr[y_j,        x_i + 1]
        z10, z11 = self.z_arr[y_j + 1, x_i], self.z_arr[y_j + 1, x_i + 1]
        # ----------------------------------------
        # 修正一：按照`双线性插值.png`的公式，z00的右边应该是z01
        # ----------------------------------------
        # z00, z10 = self.z_arr[y_j,      x_i], self.z_arr[y_j,        x_i + 1]
        # z01, z11 = self.z_arr[y_j+1, x_i], self.z_arr[y_j + 1, x_i + 1]


        # 归一化
        tx = (x - x0) / (x1 - x0)
        ty = (y - y0) / (y1 - y0)

        # 按照`双线性插值.png`的原始公式
        # zxy = 1/( (x1-x0) * (y1-y0) ) * ( z00*(x1-x)*(y1-y) + z10*(x-x0)*(y1-y) + z01*(x1-x)*(y-y0) + z11*(x-x0)*(y-y0) )

        # return z00*(1 - tx)*(1-ty) \
        #         + z10*tx*(1-ty) \
        #         + z01*(1-tx)*ty \
        #         + z11*tx*ty

        # -----------------------
        # 修正二：
        # -----------------------
        return z00*(1 - tx)*(1-ty) \
                + z01*tx*(1-ty) \
                + z10*(1-tx)*ty \
                + z11*tx*ty


# -------------------------CoordinateTransformationV7--------------------
def make_data(src, dest):
    unique_sorted_x = np.array(sorted(set(list(src[:, 0]))))
    unique_sorted_y = np.array(sorted(set(list(src[:, 1]))))
    z_x = np.empty((5, 5))
    z_y = np.empty((5, 5))

    for i in range(5):
        row = dest[i * 5:i * 5 + 5]
        for j in range(5):
            z_x[i, j], z_y[i, j] = row[j, 0], row[j, 1]

    return unique_sorted_x, unique_sorted_y, z_x, z_y


class CoordinateTransformationV7(object):
    def __init__(self, src, dest):
        (x, y, z_x, z_y) = make_data(src, dest)
        self.x_BI = xh_BI(x, y, z_x)
        self.y_BI = xh_BI(x, y, z_y)

    def get_path(self, path):

        res = []
        for i in range(len(path)):
            point = path[i]
            point_x, point_y = point[0], point[1]
            x = self.x_BI.interpolate(point_x, point_y)
            y = self.y_BI.interpolate(point_x, point_y)
            diff_x = point_x - x
            diff_y = point_y - y

            # 这里是为什么？
            res.append([point_x + diff_x, point_y + diff_y])
            # 按理说应该是下面这样
            # res.append([x, y])

        return np.array(res)







# --------------------------main---------------------------
def generate_grid(x_range: tuple | list, y_range: tuple | list, num_points: int):
    n = int(np.sqrt(num_points))
    x_points = np.linspace(x_range[0], x_range[1], n)
    y_points = np.linspace(y_range[0], y_range[1], n)

    x_points = np.tile(np.expand_dims(x_points, 0), (n, 1))
    y_points = np.tile(np.expand_dims(y_points, 1), (1, n))

    grid = np.stack([x_points, y_points], axis=-1)
    return grid.reshape((-1, 2))

def _main():
    before = np.zeros((800, 900, 3), dtype=np.uint8)
    after = np.zeros((800, 900, 3), dtype=np.uint8)

    theory = generate_grid((50, 200), (50, 200), 25)
    real = np.array([
        [50, 50], [90, 58], [125, 60], [158, 58], [200, 50],
        [56, 89], [90, 91.5], [123.6, 94.5], [160.7, 92.8], [193, 89],
        [60, 125], [92, 125], [125, 125], [157, 125], [190, 125],
        [56, 160], [90.4, 161], [125, 159.6], [160.3, 160.7], [195, 162],
        # [50, 200], [86.9, 194], [125, 190], [162, 194], [200, 200],
        [50, 200], [86.9, 194], [125, 190], [162, 194], [210, 210]
    ])

    v7 = CoordinateTransformationV7(theory, real)
    fix_theory = v7.get_path(theory)

    import cv2
    for i in range(len(real)):
        cv2.drawMarker(before, np.astype(real[i], np.int32), [0,0,255])
    for i in range(len(theory)):
        cv2.circle(before, np.astype(theory[i], np.int32), 4, [0,255,0], -1)
        cv2.circle(after, np.astype(theory[i], np.int32), 4, [0,255,0], -1)
        cv2.drawMarker(after, np.astype(fix_theory[i], np.int32), [0,0,255])

    cv2.imwrite('../output/v7_before.png', before)
    cv2.imwrite('../output/v7_after.png', after)

if __name__ == '__main__':
    _main()
    print('done.')
